// FE_solver_2D_Stokes.cpp : This file contains the 'main' function. Program execution begins and ends there.
//

#include <iostream>
#include "FEM2DVector.h"
#include <time.h>
#define PI 3.1415926535897


double nu(CPoint p, double t)
{
	return 1;
}
double  f1(CPoint p, double t)
{
	return -2 * nu(p, t)*pow(p.getX(), 2) - 2 * nu(p, t)*pow(p.getY(), 2) - nu(p, t)*exp(-p.getY()) + pow(PI, 2)*cos(PI*p.getX())*cos(2 * PI*p.getY());
}
double f2(CPoint p, double t)
{
	return 4 * nu(p, t)*p.getX()*p.getY() - nu(p, t)*pow(PI, 3)*sin(PI*p.getX()) + 2 * PI*(2 - PI * sin(PI*p.getX()))*sin(2 * PI*p.getY());
}
double u1_left(CPoint p, double t)
{
	return exp(-p.getY());
}
double u1_right(CPoint p, double t)
{
	return pow(p.getY(),2)+ exp(-p.getY());
}
double u1_bottom(CPoint p, double t)
{
	return 1.0 / 16.0*pow(p.getX(), 2) + exp(0.25);
}
double u1_up(CPoint p, double t)
{
	return 1;
}
double u2_left(CPoint p, double t)
{
	return 2;
}
double u2_right(CPoint p, double t)
{
	return -2.0 / 3.0*pow(p.getY(), 3) + 2;
}
double u2_bottom(CPoint p, double t)
{
	return 1.0 / 96.0*p.getX() + 2 - PI * sin(PI*p.getX());
}
double u2_up(CPoint p, double t)
{
	return 2 - PI * sin(PI*p.getX());
}
double fun_negative_one(CPoint, double)
{
	return -1.0;
}

double coe_fun(CPoint p, double t)
{
	return 1.0;
}

/*
double lambda(CPoint p, double t)
{
	return 1.0;
}
double mu(CPoint, double t)
{
	return 2.0;
}
double f1(CPoint p, double t)
{
	return -(lambda(p, t) + 2 * mu(p, t))*(-pow(PI, 2)*sin(PI*p.getX())*sin(PI*p.getY()))
		- (lambda(p, t) + mu(p, t))*((2 * p.getX() - 1)*(2 * p.getY() - 1)) - mu(p, t)*(-pow(PI, 2)*sin(PI*p.getX())*sin(PI*p.getY()));
}
double f2(CPoint p, double t)
{
	return -(lambda(p, t) + 2 * mu(p, t))*(2 * p.getX()*(p.getX() - 1)) -
		(lambda(p, t) + mu(p, t))*(pow(PI, 2)*cos(PI*p.getX())*cos(PI*p.getY())) - mu(p, t)*(2 * p.getY()*(p.getY() - 1));
}
double u1(CPoint p, double t)
{
	return 0;
}

double u2(CPoint p, double t)
{
	return 0;
}
double coe_fun(CPoint p, double t)
{
	return 1;
}*/

int main()
{
	/*int unknown_size = 2;
	CFEM2DVector fem(unknown_size);

	//read mesh from file
	//the mesh file contains the P,T,Pb,Tb,boundary nodes,boundary edges, and the type of the element
	fem.readMesh("..\\mesh\\cha5_exa1_qua1.neu");

	//The initializing process need mesh's information, so the function readMesh should be called first
	fem.Init();

	CBasisFunction basis_trial(TWO_DIM_QUADRATIC);
	CBasisFunction basis_test(TWO_DIM_QUADRATIC);


	fem.assignBoundNodesValue(0, 0, Dirichlet, u1);
	fem.assignBoundNodesValue(0, 1, Dirichlet, u2);


	SpMat A1 = fem.assembleMatrix2D(lambda, 0, basis_trial, 1, 0, fem.getTbTrial(), basis_test, 1, 0, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A2 = fem.assembleMatrix2D(mu, 0, basis_trial, 1, 0, fem.getTbTrial(), basis_test, 1, 0, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A3 = fem.assembleMatrix2D(mu, 0, basis_trial, 0, 1, fem.getTbTrial(),  basis_test, 0, 1, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A4 = fem.assembleMatrix2D(lambda, 0, basis_trial, 0, 1, fem.getTbTrial(),  basis_test, 1, 0, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A5 = fem.assembleMatrix2D(mu, 0, basis_trial, 1, 0, fem.getTbTrial(),  basis_test, 0, 1, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A6 = fem.assembleMatrix2D(lambda, 0, basis_trial, 1, 0, fem.getTbTrial(),  basis_test, 0, 1, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A7 = fem.assembleMatrix2D(mu, 0, basis_trial, 0, 1, fem.getTbTrial(),  basis_test, 1, 0, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});
	SpMat A8 = fem.assembleMatrix2D(lambda, 0, basis_trial, 0, 1, fem.getTbTrial(),  basis_test, 0, 1, fem.getTbTest(), pair<int, int>{fem.getPbTrial().size(), fem.getPbTest().size()});


	//SpMat A = fem.matrix_hstack(fem.matrix_vstack(A1 + 2 * A2 + A3, A6 + A7),fem.matrix_vstack(A4+A5,A8+2*A3+A2));
	SpMat A = fem.matrix_vstack(fem.matrix_hstack(A1 + 2 * A2 + A3, A4 + A5), fem.matrix_hstack(A6 + A7, A8 + 2 * A3 + A2));

	Rsv b1 = fem.assembleVector2D(f1, 0, basis_test, 0, 0, fem.getMesh().getTbTest());
	Rsv b2 = fem.assembleVector2D(f2, 0, basis_test, 0, 0, fem.getMesh().getTbTest());


	Rsv b = fem.vector_stack(b1, b2);


	fem.treatBoundaryConditions(coe_fun, 0, A, b);	


	Rsv x = fem.solveLinearEqua(DIREC_PardisoLU, A, b);

	fem.Export2VTK("result.vtk", x);
*/
			
	int unknown_size = 3;
	CFEM2DVector fem(unknown_size);
	//read mesh from file
	//the mesh file contains the P,T,Pb,Tb,boundary nodes,boundary edges, and the type of the element
	fem.readMesh("..\\mesh\\cha6_exa1_4.neu");

	//The initializing process need mesh's information, so the function readMesh should be called first
	fem.Init();

	fem.assignBoundNodesValue(0, 0, Dirichlet, u1_left);
	fem.assignBoundNodesValue(1, 0, Dirichlet, u1_bottom);
	fem.assignBoundNodesValue(2, 0, Dirichlet, u1_right);
	fem.assignBoundNodesValue(3, 0, Dirichlet, u1_up);

	fem.assignBoundNodesValue(0, 1, Dirichlet, u2_left);
	fem.assignBoundNodesValue(1, 1, Dirichlet, u2_bottom);
	fem.assignBoundNodesValue(2, 1, Dirichlet, u2_right);
	fem.assignBoundNodesValue(3, 1, Dirichlet, u2_up);
	
	CBasisFunction u_basis_trial(TWO_DIM_QUADRATIC);
	CBasisFunction u_basis_test(TWO_DIM_QUADRATIC);

	CBasisFunction p_basis_trial(TWO_DIM_LINEAR);
	CBasisFunction p_basis_test(TWO_DIM_LINEAR);

	pair<int, int> matrix_size_u{ fem.getPbTest().size(),fem.getPbTrial().size() };
	pair<int, int> matrix_size_p{ fem.getPbTest().size(),fem.getP().size() };

	//assemble matrix
	SpMat A1 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_u);
	SpMat A2 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
	SpMat A3 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);

	SpMat A5 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_p);
	SpMat A6 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_p);

	SpMat zero(fem.getP().size(), fem.getP().size());
	zero.setZero();	
	
	SpMat A = fem.matrix_vstack(fem.matrix_hstack(2 * A1 + A2, A3, A5), fem.matrix_hstack(A3.transpose(), 2 * A2 + A1, A6), fem.matrix_hstack(A5.transpose(), A6.transpose(), zero));
		
	//assemble vector
	Rsv b1 = fem.assembleVector2D(f1, 0, u_basis_test, 0, 0, fem.getTbTest());
	Rsv b2 = fem.assembleVector2D(f2, 0, u_basis_test, 0, 0, fem.getTbTest());

	Rsv zero_vec(fem.getP().size());
	zero_vec.setZero();
	
	Rsv b = fem.vector_stack(b1, b2, zero_vec);	
	
	fem.treatBoundaryConditions(coe_fun, 0, A, b, vector<unsigned int>{fem.getPbTrial().size(), fem.getPbTrial().size(), fem.getP().size()});
	//fem.ExportMat("A.txt", A);

	fem.imposeFixValue(1 + 2 * fem.getPbTest().size(), 0, A, b);
	
	
	Rsv x = fem.solveLinearEqua(DIREC_PardisoLU, A, b);
	//Rsv x = fem.solveLinearEquaParallel(A, b, 10);

	
	fem.export2VTK("u.vtk", x);

	Rsv x_p(fem.getP().size());
	for (int i=0;i<x_p.rows();++i)
	{
		x_p(i) = x(i + 2 * fem.getPbTest().size());
	}

	fem.export2VTK("p.vtk", x_p,fem.getP(),fem.getT());
}

